import pandas as pd
from MyTT import *
import akshare as ak
import numpy as np
import os
from datetime import datetime
from sqlalchemy import create_engine, distinct, or_, and_
import sqlite3
import backtrader as bt
import pymssql
from urllib.parse import quote_plus as urlquote
from configparser import ConfigParser
import warnings

warnings.filterwarnings('ignore')

"""
超级趋势
"""

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

conf = ConfigParser()
conf.read('env.ini')

sqlserver = ('sqlserver_out', 'sqlserver_in')[0]

host = conf.get(sqlserver, 'host')
port = conf.get(sqlserver, 'port')
user = conf.get(sqlserver, 'user')
password = conf.get(sqlserver, 'password')
schema = conf.get(sqlserver, 'schema')
db_name = conf.get(sqlserver, 'db')


def time_map(x):
    t = datetime.strptime(str(x), '%Y%m%d%H%M')
    return t


def time_map2(x):
    t = int(datetime.strftime(x, '%Y%m%d'))
    return t


def get_stocks_info():
    sql = "SELECT * FROM CJQS_STOCKS"
    r = exec_sql(sql)
    df2 = pd.DataFrame(r, columns=['DATE_T', 'STK_CODE'])
    df2.sort_values(by='DATE_T', inplace=True)
    df2 = df2.loc[df2['DATE_T'] >= 20230101]
    st = df2.iloc[0, 0]
    df = ak.stock_zh_index_daily(symbol="sz399300").loc[:, ['date']]
    df.rename(columns={'date': 'DATE_T'}, inplace=True)
    df['DATE_T'] = df['DATE_T'].map(time_map2)
    df = df.loc[df['DATE_T'] >= st]

    res = pd.DataFrame()

    for code in df2['STK_CODE'].unique():
        df_tmp = df2.loc[df2['STK_CODE'] == code]
        df_merge = pd.merge(df, df_tmp, how='left')
        df_merge['STK_CODE'] = df_merge['STK_CODE'].shift(1)
        df_merge.dropna(inplace=True)
        res = res.append(df_merge, ignore_index=True)

    return res


def get_data(stocks, freq='day'):
    if freq == '30min':
        table_name = 'STOCK_30MIN_DATA'
    elif freq == '60min':
        table_name = 'STOCK_60MIN_DATA'
    else:
        table_name = 'STOCK_15MIN_DATA'

    stocks_str = ','.join(f"'{tok}'" for tok in stocks)
    sql = f"SELECT * FROM {table_name} WHERE code in ({stocks_str})"
    r = exec_sql(sql)
    # print(f'table={table_name}->r={r}')
    df = pd.DataFrame(r, columns=['code', 'datetime', 'close', 'open', 'low', 'high', 'volume', 'money'])
    df['datetime'] = df['datetime'].map(time_map)
    res = dict()

    for code in stocks:
        sig_df = df.loc[df['code'] == code]
        if len(sig_df) > 0:
            sig_df.sort_values('datetime', inplace=True, ignore_index=True)
            sig_df.set_index('datetime', inplace=True)
            res[code] = sig_df
    return res


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db_name)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


def insert_db(df, table_name):
    yconnect = create_engine(f'mssql+pymssql://{user}:{urlquote(password)}@{host}:{port}/{db_name}?charset=utf8')
    pd.io.sql.to_sql(df, table_name, yconnect, schema=schema, if_exists='append', index=False)


class CJQSStrategy(bt.Strategy):
    params = (
        ('M', 26),
        ('N1', 12),
        ('N2', 144),
        ('N3', 169),
        ('N4', 576),
        ('N5', 676),
        ('M1', 60),
        ('M2', 160),
        ('P', 3),
        ('printlog', False),
    )

    def __init__(self):
        # 初始化交易指令
        self.order = None
        self.buy_price = 1e4  # 每个证券买入金额，不满一手成交一手
        self.ratio = 0.1  # 止损比例
        self.time = 3  # 止盈止损线盈亏比
        self.daily_pnl = dict()
        self.trade_info = list()
        self.buy_info = dict()  # 证券买入信息，包括买入时间和买入成本，例如{‘000001’: {'date': 20230614, 'cost': 14.23}}
        self.lb = dict()
        self.lb_count = dict()
        self.lb_fit = dict()
        self.P1 = dict()
        self.P2 = dict()
        self.P3 = dict()
        self.P4 = dict()
        self.P5 = dict()
        self.A = dict()
        self.B = dict()
        self.J0 = dict()
        self.L1 = dict()
        self.H2 = dict()
        self.buy_signal1 = dict()
        self.buy_signal2 = dict()

        for data in self.datas:
            data_name = data._name
            code = data_name[:6]
            if code not in self.buy_info:
                self.buy_info[code] = dict()

            if data_name[-2:] == 'pl':
                self.P1[code] = bt.ind.EMA(data.close, period=self.p.N1)
                self.P2[code] = bt.ind.EMA(data.close, period=self.p.N2)
                self.P3[code] = bt.ind.EMA(data.close, period=self.p.N3)
                self.P4[code] = bt.ind.EMA(data.close, period=self.p.N4)
                self.P5[code] = bt.ind.EMA(data.close, period=self.p.N5)
                self.A[code] = bt.ind.Max(self.P1[code], self.P2[code])
                self.B[code] = bt.ind.Max(self.P3[code], self.P4[code])
                self.J0[code] = (data.high + data.low + data.close + data.open) / 4.0
                self.L1[code] = bt.ind.Lowest(self.J0[code], period=self.p.M1)
                self.H2[code] = bt.ind.Highest(self.J0[code], period=self.p.M2)
                a_over_b = self.A[code] > self.B[code]
                close_over_b = data.close > self.B[code]
                # self.L1[code](-self.p.M1) 作用类似于 REF(self.L1[code], self.p.M1)
                step_up = self.L1[code] > self.L1[code](-self.p.M1)
                self.buy_signal1[code] = bt.And(a_over_b, close_over_b, step_up)

            elif data_name[-2:] == 'ps':
                vol_over_ma = data.volume > bt.ind.SMA(data.volume, period=13)
                # bt.ind.Highest(data.high(-1), period=210) 作用类似于HHV(REF(high, 1), 210)
                price_over = data.high >= bt.ind.Highest(data.high(-1), period=210)
                self.buy_signal2[code] = bt.And(vol_over_ma, price_over)

    def downcast(self, price):
        res = self.buy_price / price // 100 * 100
        res = res if res > 0 else 100
        return res

    def zt_func(self, c1, c2, ratio):
        flag = abs(c2 - c1 * ratio) < 0.01
        return flag

    def is_yizi(self, o, c, h, l):
        return o == c == h == l

    def next(self):
        if self.order:  # 查看是否有指令执行，如果有则不执行这bar
            return

        cur_dt = self.datas[0].datetime.datetime(0)
        cdt = int(cur_dt.strftime('%Y%m%d'))
        is_zt = False
        for i, data in enumerate(self.datas):
            data_name = data._name
            code = data_name[:6]
            prefix = code[:2]
            ratio = 1.2 if prefix in ('30', '68') else 1.1

            if code == '399300':
                continue
            # self.log(f'code={code}')
            try:
                dt = data.datetime.datetime(0)  # 获取当前的回测时间点
            except:
                continue

            dt_str = f'{dt.hour}:{dt.minute}'

            if data_name[-2:] == 'ps' and dt_str in ('10:30', '11:30', '14:0', '15:0'):
                if self.buy_signal1[code] and self.buy_signal2[code]:
                    self.log(f'code={code} call')

    def log(self, txt, dt=None, doprint=True):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.datetime(0)
            content = f'{dt.isoformat()}, {txt}'
            print(content)

    def notify_order(self, order):
        # 未被处理的订单
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 已经处理的订单
        if order.status in [order.Completed, order.Canceled, order.Margin]:
            code = order.data._name[:6]
            if order.isbuy():
                self.log(
                    f'code={code} BUY CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    'BUY EXECUTED, ref:%.0f，Price: %.2f, Cost: %.2f, Comm %.2f, Size: %.2f, Stock: %s' %
                    (order.ref,  # 订单编号
                     order.executed.price,  # 成交价
                     order.executed.value,  # 成交额
                     order.executed.comm,  # 佣金
                     order.executed.size,  # 成交量
                     code))  # 股票名称
                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d%H%M'))
                trade_item = [trade_dt, code, 'BUY', round(order.executed.price, 2), order.executed.size,
                              round(order.executed.comm, 2)]
                self.buy_info[code]['cost'] = order.executed.price
                self.buy_info[code]['trade_date'] = int(bt.num2date(order.created.dt).strftime('%Y%m%d'))
                self.buy_info[code]['upper_line'] = order.executed.price * (1 + self.time * self.ratio)
                self.buy_info[code]['lower_line'] = order.executed.price * (1 - self.ratio)
                self.trade_info.append(trade_item)
            else:  # Sell
                self.log(
                    f'code={code} SELL CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log('SELL EXECUTED, ref:%.0f, Price: %.2f, Cost: %.2f, Comm %.2f, Size: %.2f, Stock: %s' %
                         (order.ref,
                          order.executed.price,
                          order.executed.value,
                          order.executed.comm,
                          order.executed.size,
                          code))
                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d%H%M'))
                trade_item = [trade_dt, code, 'SELL', round(order.executed.price, 2),
                              order.executed.size, round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
                if 'last_close' in self.buy_info[code]:
                    last_close = self.buy_info[code]['last_close']
                    self.buy_info[code].clear()
                    self.buy_info[code]['last_close'] = last_close
            self.bar_executed = len(self)

        # Write down: no pending order
        self.order = None

    def stop(self):
        print('Strategy Finish!!!')

        # daily_pnl_df = pd.DataFrame(self.daily_pnl.items(), columns=['DATE_T', 'PNL'])
        # append_df = pd.DataFrame([[101, round((self.broker.getvalue() / init_cash - 1) * 100, 2)]],
        #                          columns=['DATE_T', 'PNL'])
        # daily_pnl_df = daily_pnl_df.append(append_df, ignore_index=True)
        # trade_df = pd.DataFrame(self.trade_info, columns=['DATE_T', 'STK_CODE', 'ACTION', 'PRICE', 'AMOUNT', 'COMM'])
        # pnl_table = 'TEST_DAILY_PNL'
        # trade_table = 'TEST_TRADE_INFO'
        # insert_db(daily_pnl_df, pnl_table)
        # insert_db(trade_df, trade_table)
        # print('insert table ok')


init_cash = 1e8
code_info = dict()


def run_backtrade():
    res = get_stocks_info()
    # stocks = list(res['STK_CODE'].unique())[:5]
    stocks = ['000962']
    index = '399300'
    stocks.insert(0, index)

    for dt in res['DATE_T'].unique():
        sig_df = res.loc[res['DATE_T'] == dt]
        code_info[dt] = sig_df['STK_CODE'].to_list()

    data_30min = get_data(stocks, freq='30min')
    data_60min = get_data(stocks, freq='60min')
    data0 = data_30min[index]

    cerebro = bt.Cerebro()
    datafeed0 = bt.feeds.PandasData(dataname=data0)
    cerebro.adddata(datafeed0, name=f'{index}')

    for code in data_30min:
        if code == index:
            continue
        data1 = data_30min[code]
        data2 = data_60min[code]

        datafeed1 = bt.feeds.PandasData(dataname=data1)
        cerebro.adddata(datafeed1, name=f'{code}_ps')
        datafeed2 = bt.feeds.PandasData(dataname=data2)
        cerebro.adddata(datafeed2, name=f'{code}_pl')
        print(f'{code} feeds ok')

    ##   st_date = datetime.datetime(2023, 3, 1)
    ##ed_date = datetime.datetime(2023, 6, 5)

    # 初始资金 100,000,0
    cerebro.broker.setcash(init_cash)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl')  # 返回收益率时序数据
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')  # 年化收益率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio')  # 夏普比率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')  # 回撤

    cerebro.addstrategy(CJQSStrategy)

    result = cerebro.run()
    strat = result[0]
    # 返回日度收益率序列
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # 打印评价指标
    # print("--------------- TimeReturn -----------------")
    # print(daily_return)
    # print("--------------- AnnualReturn -----------------")
    # print(strat.analyzers._AnnualReturn.get_analysis())
    # print("--------------- SharpeRatio -----------------")
    # print(strat.analyzers._SharpeRatio.get_analysis())
    # print("--------------- DrawDown -----------------")
    # print(strat.analyzers._DrawDown.get_analysis())
    # cerebro.plot()


def main():
    run_backtrade()


if __name__ == '__main__':
    main()
